use std::fmt::{self, Write};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MirType {
    I32,
    Unit,
}

impl fmt::Display for MirType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            MirType::I32 => write!(f, "i32"),
            MirType::Unit => write!(f, "()"),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Operand {
    Register(u32),
    Immediate(i32),
}

impl fmt::Display for Operand {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Operand::Register(idx) => write!(f, "_{}", idx),
            Operand::Immediate(value) => write!(f, "{}", value),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BinOpKind {
    Add,
    Sub,
    Mul,
}

impl fmt::Display for BinOpKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            BinOpKind::Add => "Add",
            BinOpKind::Sub => "Sub",
            BinOpKind::Mul => "Mul",
        };
        f.write_str(name)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CmpOpKind {
    Gt,
    Lt,
    Ge,
    Le,
    Eq,
}

impl fmt::Display for CmpOpKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let name = match self {
            CmpOpKind::Gt => "Gt",
            CmpOpKind::Lt => "Lt",
            CmpOpKind::Ge => "Ge",
            CmpOpKind::Le => "Le",
            CmpOpKind::Eq => "Eq",
        };
        f.write_str(name)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RValue {
    Binary {
        op: BinOpKind,
        left: Operand,
        right: Operand,
    },
    Compare {
        op: CmpOpKind,
        left: Operand,
        right: Operand,
    },
    Move(Operand),
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Statement {
    Assign {
        dest: u32,
        value: RValue,
    },
    SwitchInt {
        cond: Operand,
        true_target: usize,
        false_target: usize,
    },
    Goto {
        target: usize,
    },
    Return,
}

#[derive(Debug, Clone)]
pub struct BasicBlock {
    pub index: usize,
    pub statements: Vec<Statement>,
}

impl BasicBlock {
    pub fn new(index: usize) -> Self {
        Self {
            index,
            statements: Vec::new(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct MirParam {
    pub name: String,
    pub ty: MirType,
    pub register: u32,
}

#[derive(Debug, Clone)]
pub struct MirFunction {
    pub name: String,
    pub params: Vec<MirParam>,
    pub ret_ty: MirType,
    pub blocks: Vec<BasicBlock>,
}

impl MirFunction {
    pub fn new(name: impl Into<String>, ret_ty: MirType) -> Self {
        Self {
            name: name.into(),
            params: Vec::new(),
            ret_ty,
            blocks: Vec::new(),
        }
    }

    pub fn add_block(&mut self) -> usize {
        let index = self.blocks.len();
        self.blocks.push(BasicBlock::new(index));
        index
    }

    pub fn format(&self) -> String {
        let mut out = String::new();
        self.write_to(&mut out)
            .expect("formatting onto string cannot fail");
        out
    }

    pub fn write_to(&self, out: &mut String) -> fmt::Result {
        write!(out, "fn {}(", self.name)?;
        for (idx, param) in self.params.iter().enumerate() {
            if idx > 0 {
                out.push_str(", ");
            }
            write!(out, "_{}: {}", param.register, param.ty)?;
        }
        write!(out, ") -> {} {{\n", self.ret_ty)?;

        if !self.params.is_empty() {
            for param in &self.params {
                write!(out, "    debug {} => _{};\n", param.name, param.register)?;
            }
        }

        for (idx, block) in self.blocks.iter().enumerate() {
            if idx > 0 || !self.params.is_empty() {
                out.push('\n');
            }
            write!(out, "    bb{}: {{\n", block.index)?;
            for stmt in &block.statements {
                write_statement(out, stmt)?;
            }
            out.push_str("    }\n");
        }

        out.push_str("}\n");
        Ok(())
    }
}

fn write_statement(out: &mut String, stmt: &Statement) -> fmt::Result {
    out.push_str("        ");
    match stmt {
        Statement::Assign { dest, value } => {
            write!(out, "_{} = ", dest)?;
            match value {
                RValue::Binary { op, left, right } => {
                    write!(out, "{}({}, {});\n", op, left, right)?;
                }
                RValue::Compare { op, left, right } => {
                    write!(out, "{}({}, {});\n", op, left, right)?;
                }
                RValue::Move(op) => {
                    write!(out, "{};\n", op)?;
                }
            }
        }
        Statement::SwitchInt {
            cond,
            true_target,
            false_target,
        } => {
            write!(
                out,
                "switchInt({}) -> [1: bb{}, otherwise: bb{}];\n",
                cond, true_target, false_target
            )?;
        }
        Statement::Goto { target } => {
            write!(out, "goto -> bb{};\n", target)?;
        }
        Statement::Return => {
            out.push_str("return;\n");
        }
    }
    Ok(())
}

pub struct RegisterAllocator {
    next: u32,
}

impl RegisterAllocator {
    pub fn new(start: u32) -> Self {
        Self { next: start }
    }

    pub fn alloc(&mut self) -> u32 {
        let reg = self.next;
        self.next += 1;
        reg
    }
}
